import torch
import tqdm

def evaluate(model, dataloader, classes, data_conver=None):

    correct = 0
    counter = 0

    eps = 1e-6
    cls_correct = torch.tensor([0] * classes)
    cls_counter = torch.tensor([eps] * classes)

    # 统一硬件平台
    tensor_device = list(model.parameters())[-1].device
    print(tensor_device)
    cls_correct = cls_correct.to(tensor_device)
    cls_counter = cls_counter.to(tensor_device)

    pbar = tqdm.tqdm(total=len(dataloader))

    for data in dataloader:

        if data_conver == None:
            input, labels = data
        else:
            input, labels = data_conver(data)

        with torch.no_grad():
            input = input.to(tensor_device)
            labels = labels.to(tensor_device)
            pred = model(input)
            results = torch.max(pred, 1)[1]

            correct += (results == labels).sum().item()
            counter += labels.shape[0]

            mask = torch.tensor([list(range(classes))] * labels.shape[0], device=tensor_device).transpose(0, 1)
            cls_counter += (mask == labels).sum(dim=1)
            # classes+1代表没预测对
            cls_correct += (torch.where(results == labels, labels, classes+1) == mask).sum(dim=1)
            pbar.update(1)

    rec = correct / counter
    rec_cls = cls_correct / cls_counter
    rec_cls = rec_cls.tolist()

    return rec, rec_cls

# 将dict转换成feature和label的形式
# 用于转换dataloader的输出形式
# 以使evaluate函数适用于更多形式的数据集
def dict2ftlb(dict):
    feature = dict['data']
    label = dict['fine_labels']
    return feature, label

# def model(a):
#     ids = torch.randint(5, size=(5, ))
#     result = torch.randint(5, size=(5, 5))
#     result[list(range(5)), ids] = 7
#     # print(ids)
#     # print(result)
#     # print("=====")
#     return result

if __name__ == "__main__":
    from options import opt
    from CIFAR import CIFAR
    import torchvision.transforms as trans
    from torch.utils.data import DataLoader
    from networks import Model
    from logger import logger
    val_dataset = CIFAR(
        {
            'folder'    : '../CIFAR/CIFAR100',
            'name'      : 'test',
            'list_file' : 'test_list.txt'
        },
        transforms=trans.Compose([
            trans.ToTensor(),
            trans.Resize(224)
        ])
    )
    val_loader = DataLoader(
        val_dataset, 
        batch_size=opt.batch, 
        num_workers=opt.workers, 
        shuffle=False
    )
    model = Model(opt)
    if opt.load:
        load_epoch = model.load(opt.load)
    log_file = logger(opt)
    rec, rec_cls = evaluate(model.classifier, val_loader, 100, dict2ftlb)
    log_file.append({
                'rec' : rec, 
                'rec_cls' : str(rec_cls),
            })